# -*- coding: utf-8 -*-
"""
@File: kline.py
@author: chenxialin
@time: 9月 26, 2025 21:35
"""
# api/kline.py
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel
from sqlalchemy import or_
from sqlalchemy.orm import Session
from typing import Optional, List, Dict, Any
import requests
import pandas as pd

from mysql.database import SessionLocal, Base
from models.kline import Kline  # 你的模型



# 依赖项：获取 DB 会话
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

class KlineRequest(BaseModel):
    code: str
    # fqt: str = "1"  # 复权类型："0":不复权, "1":前复权, "2":后复权
    # timeout: int = 600000  # 单位：毫秒

def get_qt_klines(secid: str, fqt: str = "1", timeout_ms: int = 10000) -> Optional[Dict[str, Any]]:
    url = "https://push2his.eastmoney.com/api/qt/stock/kline/get"
    timeout_sec = timeout_ms / 1000.0

    params = {
        "secid": secid,
        "fields1": "f1,f2,f3,f4,f5,f6",
        "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61",
        "klt": "101",
        "fqt": fqt,
        "end": "20500101",
        "lmt": "1000000",
    }

    headers = {
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) "
            "Chrome/120.0.0.0 Safari/537.36"
        ),
        "Referer": f"http://quote.eastmoney.com/{secid.split('.')[-1]}.html",
    }
   

    try:
        response = requests.get(url, params=params, headers=headers, timeout=timeout_sec)
        response.raise_for_status()
        data = response.json()
        return data.get("data")
    except Exception as e:
        print(f"请求失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"请求失败: {str(e)}")


EastMoneyMarketMap={
    "SH":"1",
    "SZ":"0",
    "BJ":"0",
}

def fetch_and_save_kline(data: KlineRequest):
    db_gen = get_db()
    db: Session = next(db_gen)
    existing_kline = db.query(Kline).filter(Kline.code == data.code).first()
    if existing_kline:
        return {"message": "该股票的K线数据已存在"}

    secid=f"{EastMoneyMarketMap[data.code.split('.')[1]]}.{data.code.split('.')[0]}"
    raw_data = get_qt_klines(secid=secid, fqt="0", timeout_ms=600000)
    if not raw_data:
        raise HTTPException(status_code=404, detail="未获取到K线数据，请检查 secid")

    code = data.code
    klines = raw_data.get("klines", [])
    if not klines:
        # raise HTTPException(status_code=404, detail="该股票无K线数据")
        return {"message": "该股票无K线数据"}

    # 检查数据库中是否已存在该股票的K线数据
   

    # 获取股票代码
    code = data.code

    # 解析数据
    columns = ["date", "open", "close", "high", "low", "volume", "amount", "amplitude", "chg", "change", "turn"]
    df = pd.DataFrame([item.split(",") for item in klines], columns=columns)

    # 类型转换
    df["date"] = df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y%m%d").astype(int)  # 转为 20250926 格式
    df["open"] = pd.to_numeric(df["open"], errors='coerce')
    df["close"] = pd.to_numeric(df["close"], errors='coerce')
    df["high"] = pd.to_numeric(df["high"], errors='coerce')
    df["low"] = pd.to_numeric(df["low"], errors='coerce')
    df["volume"] = pd.to_numeric(df["volume"], errors='coerce')
    df["amount"] = pd.to_numeric(df["amount"], errors='coerce')
    df["amplitude"] = pd.to_numeric(df["amplitude"], errors='coerce')
    df["chg"] = pd.to_numeric(df["chg"], errors='coerce')
    df["change"] = pd.to_numeric(df["change"], errors='coerce')
    df["turn"] = pd.to_numeric(df["turn"], errors='coerce')

    # 过滤空值
    df.dropna(inplace=True)

    # 构造模型对象
    new_records = []
    existing_dates = []

    # 查询已存在的 (code, date) 组合，避免重复插入
    existing_date_set = set(
        db.query(Kline.date).filter(Kline.code == code).all()
    )
    existing_dates = {row[0] for row in existing_date_set}

    for _, row in df.iterrows():
        date_int = int(row['date'])
        if date_int in existing_dates:
            continue  # 跳过已存在的记录

        record = Kline(
            code=code,
            date=date_int,
            open=row['open'],
            close=row['close'],
            high=row['high'],
            low=row['low'],
            volume=row['volume'],
            amount=row['amount'],
            amplitude=row['amplitude'],
            chg=row['chg'],
            change=row['change'],
            turn=row['turn']
        )
        new_records.append(record)

    if not new_records:
        print(f"{code}无新K线数据需要插入")
        return {"message": "无新数据需要插入"}

    # 批量插入
    try:
        db.bulk_save_objects(new_records)
        db.commit()
        print(f"{code}成功插入 {len(new_records)} 条K线数据")
        return {
            "message": f"成功插入 {len(new_records)} 条K线数据",
            "code": code,
            "inserted": len(new_records)
        }
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=f"数据库插入失败: {str(e)}")
